/*
  Glues a imagem sequence passed as argument and generate a new image 
  with all matched images glued.
  
  This file is based on match.c.
 
  Copyright (C) 2006-2007  Rob Hess <hess@eecs.oregonstate.edu>
 
  @version 1.1.1-20070913
*/

#include <syslog.h>
#include <sys/time.h>
#include <time.h>
#include "egdebug.h"
#include "sift.h"
#include "imgfeatures.h"
#include "kdtree.h"
#include "utils.h"
#include "xform.h"
#include "gluesh.h"

#include <cv.h>
#include <cxcore.h>
#include <highgui.h>

#include <stdio.h>

#define MAXSTR 128

typedef enum {
	FINDFEATURES=0,
	GLUEIMAGE,
	SIFT1,
	SIFT2,
	KDTREE,
	RANSAC,
	MAX_TIME_TABLE_INDEX
} timetableindex_t;

long int timetable[MAX_TIME_TABLE_INDEX];

char timetablestr[MAX_TIME_TABLE_INDEX][MAXSTR] =
{
	"FindFeatures",
	"GlueImage",
	"Sift1",
	"Sift2",
	"KdTree",
	"Ransac"
};

void showTimeTable()
{
	int i;
	for (i=0; i<MAX_TIME_TABLE_INDEX; i++)
	{
		fprintf(stderr,"%s = \t%ld\n",timetablestr[i],timetable[i]);
	}
}

void initTimeTable()
{
	int i;
	for (i=0; i<MAX_TIME_TABLE_INDEX; i++)
	{
		timetable[i] = 0;
	}
}

int main( int argc, char** argv )
{
    IplImage *img2, *newimg;
    feature_t *feat1, *feat2;
    int i, m, n1;
    struct timeval t1;
    
	EG_INI();
	EG_INF("Starting e-glue.");
	
	initTimeTable();
	
    if( argc < 3 )
    {
        EG_ERR( "usage: %s <img1> ... <imgn>", argv[0] );
        exit(1);
    }
    
    newimg = cvLoadImage( argv[1], 1 );
    if( ! newimg )
    {
        EG_ERR( "unable to load image from %s", argv[1] );
        exit(1);
    }
    printf("Processing image %s\n", argv[1]);
    for (i=2; i<argc; i++)
    {
	    img2 = cvLoadImage( argv[i], 1 );
	    if( ! img2 )
	    {
	        EG_ERR( "unable to load image from %s", argv[i] );
	        exit(1);
	    }    	
	    tic(&t1);
		m = egFindFeatures(newimg, &feat1, &n1, img2, &feat2);
		timetable[FINDFEATURES] += tac(t1);
		if (m < EG_MIN_SIFT_MATCHES)
		{
			EG_INF("Matches found: %d\n", m);
			printf("Discarding image %s\n", argv[i]);
			cvReleaseImage(&img2);
			continue;
		}
    	IplImage *tmpimg;
		printf("Processing image %s\n", argv[i]);
		tic(&t1);
    	if (egGlue(newimg, feat1, n1, img2, feat2, &tmpimg))
    	{
    		timetable[GLUEIMAGE] += tac(t1);
    		cvReleaseImage(&img2);	
    		cvReleaseImage(&newimg);
    		newimg = tmpimg;	
    	}
    	showTimeTable();
    }
    if ( newimg )
    {
    	cvSaveImage("newimg.png", newimg);
    	printf("Image saved in newimg.png\n");
    	cvReleaseImage( &newimg );
    }
    cvReleaseImage( &img2 );
    free( feat1 );
    free( feat2 );
	EG_INF("Finishing e-glue.");    
    EG_END();
    return 0;
}

void PrintMat(CvMat *A)
{
    int i, j, tmp_i;
    char tmp[256], tmp2[256];
    EG_DEB( "Entering %s", __FUNCTION__);
    for (i = 0; i < A->rows; i++)
    {
        switch (CV_MAT_DEPTH(A->type))
        {
        case CV_32F:
        case CV_64F:
        	memset(tmp, 0, 256);
        	memset(tmp2, 0, 256);
            for (j = 0; j < A->cols; j++)
            {
                snprintf(tmp2, 256, "%8.3f ", (float)cvGetReal2D(A, i, j));
                tmp_i = strlen(tmp);
                strncat(tmp + tmp_i, tmp2, 256 - tmp_i - strlen(tmp2));
            }
            EG_INF(tmp);
            break;
        case CV_8U:
            case CV_16U:
		        	memset(tmp, 0, 256);
		        	memset(tmp2, 0, 256);
                    for(j = 0; j < A->cols; j++)
                    {
                        snprintf(tmp2, 256, "%6d", (int)cvGetReal2D(A, i, j));
		                tmp_i = strlen(tmp);
		                strncat(tmp + tmp_i, tmp2, 256 - tmp_i - strlen(tmp2));
                    }
                    EG_INF(tmp);
            break;
        default:
                break;
        }
    }
    EG_DEB("Exiting %s", __FUNCTION__);
}

void egGetSizeDisp(CvSize size1, CvSize size2, CvMat* H_inv, int *width, int *height, int *dispx, int *dispy)
{
	int i;
	int minx, miny, maxx, maxy;
    EG_DEB( "Entering %s", __FUNCTION__);
	CvMat *new_dim = cvCreateMat(3, 4, H_inv->type);
	egGuessNewDim(size2, H_inv, new_dim);
	EG_DEB("\nnew_dim:");
	PrintMat(new_dim);
	minx = 0;
	miny = 0;
	maxx = size1.width; 
	maxy = size1.height; 
	for (i=1; i < new_dim->cols; i++)
	{
		if (minx > cvmGet(new_dim, 0, i))
		{
			minx = cvmGet(new_dim, 0, i);
		}
		if (maxx < cvmGet(new_dim, 0, i))
		{
			maxx = cvmGet(new_dim, 0, i);
		}
		if (miny > cvmGet(new_dim, 1, i))
		{
			miny = cvmGet(new_dim, 1, i);
		}
		if (maxy < cvmGet(new_dim, 1, i))
		{
			maxy = cvmGet(new_dim, 1, i);
		}
	}
	*height = maxy - miny;
	*width = maxx - minx;
	*dispx = minx;
	*dispy = miny;
	cvReleaseMat( &new_dim );
	EG_DEB("height = %d\n", *height);
	EG_DEB("width = %d\n", *width);
	EG_DEB("dispx = %d\n", *dispx);
	EG_DEB("dispy = %d\n", *dispy);
    EG_DEB( "Exiting %s", __FUNCTION__);
}

void egNorm(CvMat* src, CvMat* dst)
{
	int i;
    EG_DEB( "Entering %s", __FUNCTION__);
	assert(src->rows == 3);
	for (i=0; i < src->cols; i++)
	{
		double factor = cvmGet(src, 2, i);
		cvmSet(dst, 0, i, cvmGet(src, 0, i) / factor);
		cvmSet(dst, 1, i, cvmGet(src, 1, i) / factor);
		cvmSet(dst, 2, i, 1.0);
	}
    EG_DEB( "Exiting %s", __FUNCTION__);
}

void egNorm3x3(CvMat* src, CvMat* dst)
{
	int i,j;
	double factor;
    EG_DEB( "Entering %s", __FUNCTION__);
	EG_DEB("rows = %d\n cols = %d", src->rows, src->cols); 
	assert(src->rows == 3 && src->cols == 3);
	factor = cvmGet(src, 2, 2);
	/* TODO: look for a matrix to scalar multiplication. */
	for (i=0; i < src->cols; i++)
	{
		for (j=0; j < src->rows; j++)
		{
			cvmSet(dst, j, i, cvmGet(src, j, i) / factor);
		}
	}
    EG_DEB( "Exiting %s", __FUNCTION__);
}

void egGuessNewDim(CvSize size, CvMat* H_inv, CvMat* new_dim)
{
    EG_DEB( "Entering %s", __FUNCTION__);
	CvMat *dim = cvCreateMat(3, 4, H_inv->type);
	cvmSet(dim, 0, 0, 0);
	cvmSet(dim, 1, 0, 0);
	cvmSet(dim, 2, 0, 1);
	cvmSet(dim, 0, 1, size.width);
	cvmSet(dim, 1, 1, 0);
	cvmSet(dim, 2, 1, 1);
	cvmSet(dim, 0, 2, 0);
	cvmSet(dim, 1, 2, size.height);
	cvmSet(dim, 2, 2, 1);
	cvmSet(dim, 0, 3, size.width);
	cvmSet(dim, 1, 3, size.height);
	cvmSet(dim, 2, 3, 1);
	cvMatMul(H_inv, dim, new_dim);
	egNorm(new_dim, new_dim);
    EG_DEB( "Exiting %s", __FUNCTION__);
}

void egBuildH_disp(int dispx, int dispy, CvMat* H_disp)
{
    EG_DEB( "Entering %s", __FUNCTION__);
	cvmSet(H_disp, 0, 0, 1.0);
	cvmSet(H_disp, 0, 1, 0);
	cvmSet(H_disp, 0, 2, -dispx);
	cvmSet(H_disp, 1, 0, 0);
	cvmSet(H_disp, 1, 1, 1.0);
	cvmSet(H_disp, 1, 2, -dispy);
	cvmSet(H_disp, 2, 0, 0);
	cvmSet(H_disp, 2, 1, 0);
	cvmSet(H_disp, 2, 2, 1.0);
    EG_DEB( "Exiting %s", __FUNCTION__);
}

void egBuildH_trans(CvMat* H_inv, CvMat* H_disp, CvMat* H_trans)
{
    EG_DEB( "Entering %s", __FUNCTION__);
	cvMatMul(H_inv, H_disp, H_trans);
	egNorm3x3(H_trans, H_trans);
    EG_DEB( "Exiting %s", __FUNCTION__);
}

/**
 * Find features based on two input images.
 *  
 * Note that this line:
 *      
 * feat1[i].fwd_match = nbrs[0];
 *      
 * is important for the RANSAC function to work.
 * 
 */
int egFindFeatures(IplImage *img1, feature_t** feat1, int *n1, IplImage *img2, feature_t **feat2)
{
	IplImage *stacked;
    feature_t *feat;
	struct kd_node* kd_root;
    double d0, d1;
	feature_t **nbrs;
    CvPoint pt1, pt2;
	int i, k, m = 0;
	int n2;
	struct timeval t1;
    EG_DEB( "Entering %s", __FUNCTION__);

	stacked = stack_imgs( img1, img2 );
	tic(&t1);
	*n1 = _sift_features( img1, feat1, SIFT_INTVLS, SIFT_SIGMA, EG_SIFT_CONTR_THR,
			 EG_SIFT_CURV_THR, SIFT_IMG_DBL, SIFT_DESCR_WIDTH,
			 SIFT_DESCR_HIST_BINS );
	timetable[SIFT1] += tac(t1);
	tic(&t1);
	n2 = _sift_features( img2, feat2, SIFT_INTVLS, SIFT_SIGMA, EG_SIFT_CONTR_THR,
			 EG_SIFT_CURV_THR, SIFT_IMG_DBL, SIFT_DESCR_WIDTH,
			 SIFT_DESCR_HIST_BINS );
 	timetable[SIFT2] += tac(t1);
//	*n1 = sift_features( img1, feat1 );
//   n2 = sift_features( img2, feat2 );
	tic(&t1);
    kd_root = kdtree_build( *feat2, n2 );
    for( i = 0; i < *n1; i++ )
    {
        feat = *feat1 + i;
        k = kdtree_bbf_knn( kd_root, feat, 2, &nbrs, KDTREE_BBF_MAX_NN_CHKS );
        if( k == 2 )
        {
            d0 = descr_dist_sq( feat, nbrs[0] );
            d1 = descr_dist_sq( feat, nbrs[1] );
            if( d0 < d1 * NN_SQ_DIST_RATIO_THR )
            {
                pt1 = cvPoint( cvRound( feat->x ), cvRound( feat->y ) );
                pt2 = cvPoint( cvRound( nbrs[0]->x ), cvRound( nbrs[0]->y ) );
                pt2.y += img1->height;
                cvLine( stacked, pt1, pt2, CV_RGB(255,0,255), 1, 8, 0 );
                m++;
                (*feat1)[i].fwd_match = nbrs[0];
            }
        }
        free( nbrs );
    }
	timetable[KDTREE] += tac(t1);
    EG_DEB("Found %d total matches\n", m );
//    display_big_img( stacked, "Matches" );
//    cvSaveImage("stacked.png", stacked);
//    cvWaitKey( 0 );
    cvReleaseImage( &stacked );
    EG_DEB("After release stacked\n");
    kdtree_release( kd_root );
    EG_DEB( "Exiting %s", __FUNCTION__);
    return m;
}

int egGlue(IplImage *img1, feature_t *feat1, int n1, IplImage *img2, feature_t *feat2, IplImage **newimg)
{
    IplImage* xformed2, *xformed1;
   	int dispx, dispy; /* x and y displacement */
    CvMat *H, *H_inv, *H_disp, *H_trans;
    IplImage* mask;
    struct timeval t1;
//    static int toblend = 0;
//    char toblend_name[256];
    CvSize xformed_size;
    EG_DEB( "Entering %s", __FUNCTION__);
    tic(&t1);
    H = ransac_xform( feat1, n1, FEATURE_FWD_MATCH, lsq_homog, EG_MIN_RANSAC_MATCHES, 0.01,
                      homog_xfer_err, 3.0, NULL, NULL );
    timetable[RANSAC] += tac(t1);
    EG_DEB("After RANSAC.\n");
    if( H )
    {
        H_inv = cvCreateMat(3, 3, H->type);
  		EG_DEB("After cvClone.\n");
        if( H_inv )
        {
        	H_disp = cvCreateMat(3, 3, H->type);
        	H_trans = cvCreateMat(3, 3, H->type);
    		EG_DEB("After cvCreateMat.\n");
        	/* computes the image size after transform. */
			EG_DEB("\nH:");
			PrintMat(H);
		    if ( cvInvert(H, H_inv, CV_SVD) == 0)
		    {
		    	EG_ERR("Can't invert transform.\n");
	            cvReleaseMat( &H );
	            cvReleaseMat( &H_inv );
	            cvReleaseMat( &H_disp );
	            cvReleaseMat( &H_trans );
			    EG_DEB( "Exiting 1 %s", __FUNCTION__);
		    	return 0;
		    }
//			egNorm3x3(H_inv, H_inv);
			EG_DEB("\nH_inv:");
			PrintMat(H_inv);
        	egGetSizeDisp(cvGetSize( img1 ), cvGetSize( img2 ), H_inv, 
        				  &xformed_size.width, &xformed_size.height, &dispx, &dispy);
        	/* resets the displacement the fit all the new image*/
			egBuildH_disp(dispx, dispy, H_disp);
			egBuildH_trans(H_inv, H_disp, H_trans); 
			EG_DEB("\nH_trans:");
			PrintMat(H_trans);
            xformed2 = cvCreateImage( xformed_size, IPL_DEPTH_8U, 3 );
            xformed1 = cvCreateImage( xformed_size, IPL_DEPTH_8U, 3 );
            *newimg = cvCreateImage( xformed_size, IPL_DEPTH_8U, 3 );
            /* Ensure that the both images will be at the same coordinates system */
            cvWarpPerspective( img2, xformed2, H_trans,
                               CV_INTER_LINEAR + CV_WARP_FILL_OUTLIERS,
                               cvScalarAll( 0 ) );
//           	snprintf(toblend_name,256,"toblend_%03d.tif", toblend++);
//           	cvSaveImage(toblend_name, xformed2);
            EG_DEB("After img2\n");
            cvWarpPerspective( img1, xformed1, H_disp,
                               CV_INTER_LINEAR + CV_WARP_FILL_OUTLIERS,
                               cvScalarAll( 0 ) );
            EG_DEB("After img1\n");
           mask = egComputeMask(xformed1, xformed2);
           cvCopy( xformed1, *newimg, NULL );
           cvCopy( xformed2, *newimg, mask ); /* replaces only some pixels */
////            cvNamedWindow( "Xformed2", 1 );
////            cvShowImage( "Xformed", xformed2 );
////            cvNamedWindow( "Xformed1", 1 );
////            cvShowImage( "Xformed1", xformed1 );
////            cvWaitKey( 0 );
////            cvSaveImage( "xformed1.png", xformed1 );
////            cvSaveImage( "xformed2.png", xformed2 );
//            cvCopy( xformed2, *newimg, NULL );
//            CvRect rect = cvRect(0, 0, img1->width, img1->height);
//            cvSetImageROI( xformed1, rect );
//            cvSetImageROI( *newimg, rect );
//            cvCopy( xformed1, *newimg, NULL );
//            cvResetImageROI( *newimg );
//            cvSaveImage( "newimg.png", *newimg );
            cvReleaseImage( &xformed1 );
            cvReleaseImage( &xformed2 );
            cvReleaseMat( &H_inv );
            cvReleaseMat( &H_disp );
            cvReleaseMat( &H_trans );
        }
        cvReleaseMat( &H );
	    EG_DEB( "Exiting 2 %s", __FUNCTION__);
        return 1;
    }
    else
    {
    	EG_DEB("Can't compute RANSAC.\n");
	    EG_DEB( "Exiting 3 %s", __FUNCTION__);
    	return 0;
    }    
}

IplImage* egComputeMask(IplImage *xformed1, IplImage *xformed2)
{
	IplImage *tmp1, *tmp2;
    EG_DEB( "Entering %s", __FUNCTION__);
	IplImage *mask = cvCreateImage( cvGetSize(xformed1), IPL_DEPTH_8U, 1);
	IplImage *tmp1arr = cvCreateImage( cvGetSize(xformed1), IPL_DEPTH_8U, 1);
	IplImage *tmp2arr = cvCreateImage( cvGetSize(xformed2), IPL_DEPTH_8U, 1);
	tmp1 = cvCreateImage( cvGetSize(xformed1), IPL_DEPTH_8U, 3 );
	tmp2 = cvCreateImage( cvGetSize(xformed2), IPL_DEPTH_8U, 3 );
	cvCvtColor( xformed1, tmp1, CV_RGB2YCrCb);
	cvCvtColor( xformed2, tmp2, CV_RGB2YCrCb);
	cvSplit(tmp1, tmp1arr, NULL, NULL, NULL);
	cvSplit(tmp2, tmp2arr, NULL, NULL, NULL);
	cvCmp( tmp1arr, tmp2arr, mask, CV_CMP_LT );
	cvReleaseImage( &tmp1 ); 
	cvReleaseImage( &tmp2 );
	cvReleaseImage( &tmp1arr );
	cvReleaseImage( &tmp2arr );
    EG_DEB( "Exiting %s", __FUNCTION__);
	return mask;
}
